import os
import mmap
import collections
from tqdm import tqdm

def iter_lines(filename, chunk_size=1024*1024):
    """
    按块读取文件，并保证每次返回完整的行，不会一次将整个文件加载到内存中。
    
    参数：
      filename   : 文件路径
      chunk_size : 每次映射的字节数（可根据实际情况调优）
    
    产生：
      每一行文本（已解码为 UTF-8）
    """
    with open(filename, 'rb') as f:
        file_size = os.fstat(f.fileno()).st_size
        offset = 0
        leftover = b""
        while offset < file_size:
            # 每次映射固定大小的块
            bytes_to_map = min(chunk_size, file_size - offset)
            mm = mmap.mmap(f.fileno(), length=bytes_to_map, offset=offset, access=mmap.ACCESS_READ)
            # 将前一次块末尾残留的数据拼接上本次映射的数据
            data = leftover + mm[:]
            mm.close()
            # 找到最后一个换行符，保证本块处理完整行
            last_newline = data.rfind(b'\n')
            if last_newline != -1:
                # 分割出完整行
                chunk_lines = data[:last_newline].split(b'\n')
                for line in chunk_lines:
                    yield line.decode('utf-8')
                # 将剩余未完成的一行保留，待下次块继续处理
                leftover = data[last_newline+1:]
            else:
                # 如果本块内没有换行符，则将数据全部保留
                leftover = data
            offset += bytes_to_map
        # 如果文件末尾还有残留数据，则作为最后一行返回
        if leftover:
            yield leftover.decode('utf-8')

def count_frequency(filename, process_fn=lambda line: line.strip(), ignore_one=False):
    """
    统计超大文件中每行 key 的频率，每行数据视为一个 key。
    
    参数：
      filename : 文件路径
    返回：
      collections.Counter 对象，保存每个 key 的出现次数
    """
    counter = collections.Counter()
    with tqdm(desc="Counting Frequencies") as pbar:
        for line in iter_lines(filename):
            key = process_fn(line)
            if key:
                counter[key] += 1
            pbar.update(1)
    return counter if not ignore_one else {k: v for k, v in counter.items() if v > 1}

if __name__ == "__main__":
    filename = "./datas/valid/split/documents_2_multithread_question_type_name.jsonl"  # 请替换为你的文件路径
    freq = count_frequency(filename)
    # 输出频率最高的 key
    for key, count in freq.most_common():
        print(f"{key}: {count}")
